//  Copyright 2011 Charlie Garrett
//
//  Licensed under the Apache License, Version 2.0 (the "License");
//  you may not use this file except in compliance with the License.
//  You may obtain a copy of the License at
//
//  http://www.apache.org/licenses/LICENSE-2.0
//
//  Unless required by applicable law or agreed to in writing, software
//  distributed under the License is distributed on an "AS IS" BASIS,
//  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
//  See the License for the specific language governing permissions and
//  limitations under the License.
//
//  CodeGenerator.cpp
//

#include "CodeGenerator.h"

namespace logic {
    using namespace std;

    void CodeGenerator::InitState() {
        ClearRegisters();
        
        entry_points.clear();
    }
    
    void CodeGenerator::ClearRegisters() {
        ClearTempRegisters();
        ClearPermanentRegisters();
        ClearArgRegisters();
    }       
        
    void CodeGenerator::ClearTempRegisters() {
        next_temp_reg = 0;
        temp_regs.clear();
        temp_available_term.clear();
    }
    
    void CodeGenerator::ClearPermanentRegisters() {
        next_permanent_reg = 0;
        permanent_regs.clear();
        permanent_available_term.clear();
    }
    
    void CodeGenerator::ClearArgRegisters() {
        next_arg_reg = 0;
        arg_regs.clear();
        arg_available_term.clear();
    }
    
    Register CodeGenerator::AllocateRegister(Register::Kind kind, Term *term) {
        Register reg;
        switch (kind) {
            case Register::UNALLOCATED:
                return Register::Unallocated;
            case Register::TEMPORARY:
                assert(temp_regs.find(term) == temp_regs.end());
                reg = Register(Register::TEMPORARY, next_temp_reg++);
                temp_regs[term] = reg;
                return reg;
            case Register::PERMANENT:
                assert(permanent_regs.find(term) == permanent_regs.end());
                reg = Register(Register::PERMANENT, next_permanent_reg++);
                permanent_regs[term] = reg;
                return reg;
            case Register::ARGUMENT:
                assert(arg_regs.find(term) == arg_regs.end());
                reg = Register(Register::ARGUMENT, next_arg_reg++);
                arg_regs[term] = reg;
                return reg;                
        }
    }
    
    bool CodeGenerator::IsRegisterUsed(Register reg) {
        switch (reg.kind) {
            case Register::UNALLOCATED:
                return false;
            case Register::TEMPORARY:
                return temp_available_term.find(reg) != temp_available_term.end();
            case Register::PERMANENT:
                return permanent_available_term.find(reg) != permanent_available_term.end();
            case Register::ARGUMENT:
                return arg_available_term.find(reg) != arg_available_term.end();
        }
    }
    
    // Get a register of the given kind previously allocated to the term, if one exists.
    Register CodeGenerator::GetRegister(Register::Kind kind, Term *term) {
        switch (kind) {
            case Register::UNALLOCATED:
                return Register::Unallocated;
            case Register::TEMPORARY:
                if (temp_regs.find(term) != temp_regs.end()) {
                    return temp_regs[term];
                }
                break;
            case Register::PERMANENT:
                if (permanent_regs.find(term) != permanent_regs.end()) {
                    return permanent_regs[term];
                }
                break;
            case Register::ARGUMENT:
                if (arg_regs.find(term) != arg_regs.end()) {
                    return arg_regs[term];
                }
        }
        return Register::Unallocated;
    }

    void CodeGenerator::EmitInstruction(Instruction *inst, Term *term) {
        // Add label to the instruction if there is one
        if (valid_label) {
            inst->SetLabel(next_instruction_label);
            valid_label = false;
        }
        interpreter->EmitInstruction(inst);
        
        if (term != NULL) {
            // TODO: This should only set the registers written by the instruction.
            Register reg = GetRegister(Register::TEMPORARY, term);
            if (reg.kind != Register::UNALLOCATED) {
                temp_available_term[reg] = term;
            }
            reg = GetRegister(Register::PERMANENT, term);
            if (reg.kind != Register::UNALLOCATED) {
                permanent_available_term[reg] = term;
            }
            reg = GetRegister(Register::ARGUMENT, term);
            if (reg.kind != Register::UNALLOCATED) {
                arg_available_term[reg] = term;
            }
        }
    }
    
    Instruction::Label CodeGenerator::LabelInstruction(const Functor &functor,
                                                         size_t clause_num) {
        Instruction::Label key(functor, clause_num);
        entry_points[key] = interpreter->state.code.size();
        next_instruction_label = key;
        valid_label = true;
        return key;
    }
    
    bool CodeGenerator::LabelOffset(const Instruction::Label &key, size_t *offset) {
        if (entry_points.find(key) == entry_points.end()) {
            *offset = 0;
            return false;
        } else {
            *offset = entry_points[key];
            return true;
        }
    }
    
    bool CodeGenerator::ResolveLabels() {
        for (vector<Instruction*>::iterator inst = interpreter->state.code.begin();
             inst != interpreter->state.code.end(); ++inst) {
            if (!(*inst)->ResolveLabel(this)) {
                cout << "Failed to resolve label for: " << (*inst)->ToString() << endl;
                return false;
            }
        }
        return true;
    }
    
    void CodeGenerator::RemoveInstructions(Instruction::Label start) {
        map<Instruction::Label, size_t>::iterator pos = entry_points.find(start);
        assert(pos != entry_points.end());

        vector<Instruction*>::iterator inst_start =
            interpreter->state.code.begin() + pos->second;
        vector<Instruction*>::iterator inst_end = inst_start + 1;
        for (; inst_end != interpreter->state.code.end(); ++inst_end) {
            if ((*inst_end)->HasLabel()) {
                break;
            }
        }
        interpreter->state.code.erase(inst_start, inst_end);
        entry_points.erase(pos);
    }
    
}  // namespace logic
